package org.huzhp.gateway.filter.balance;

import com.netflix.loadbalancer.Server;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.DelegatingServiceInstance;
import org.springframework.cloud.gateway.support.NotFoundException;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.Ordered;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.util.Assert;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.util.UriComponentsBuilder;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * @Description: 负载均衡拦截器
 * @Author: huzp
 * @CreateDate: 2020/4/15 16:21
 * @Version: 1.0
 */
public class LoadBalancerGatewayFilter implements GatewayFilter, Ordered {

    private static final Logger logger = LoggerFactory.getLogger(LoadBalancerGatewayFilter.class);

    private GatewayLoadBalanceClient loadBalanceClient;

    private String rule;

    private List<Server> servers = new ArrayList<>();


    public LoadBalancerGatewayFilter(GatewayLoadBalanceClient gatewayLoadBalanceClient,String rule,List<String> servers) {
        this.loadBalanceClient =  gatewayLoadBalanceClient;
        this.rule = rule;
        for (String server : servers){
            if (LbRule.Weight.name().equalsIgnoreCase(rule)){
                String[] serverArr = server.split("\\s+");
                this.servers.add(createServer(serverArr[0],serverArr.length==2?serverArr[1]:null));
            }else{
                this.servers.add(createServer(server));
            }

        }
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        URI url = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR);
        Route route = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
        String schemePrefix = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_SCHEME_PREFIX_ATTR);
        if (url != null && ("lb".equals(url.getScheme()) || "lb".equals(schemePrefix))) {
            ServiceInstance instance;
            ServerWebExchangeUtils.addOriginalRequestUrl(exchange, url);

            instance = loadBalanceClient.choose(route.getId(),rule,servers,getHint(exchange.getRequest()));

            if (instance == null) {
                throw new NotFoundException( "Unable to find instance for " + url.getHost());
            } else {
                URI uri = exchange.getRequest().getURI();
                String overrideScheme = instance.isSecure() ? "https" : "http";
                if (schemePrefix != null) {
                    overrideScheme = url.getScheme();
                }
                logger.info("load balance {}",instance.getUri());
                URI requestUrl = UriComponentsBuilder.fromUri(uri).scheme(overrideScheme).host(instance.getHost()).port(instance.getPort()).build().toUri();

                exchange.getAttributes().put(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR, requestUrl);
                return chain.filter(exchange);
            }
        } else {
            return chain.filter(exchange);
        }
    }

    private String getHint(ServerHttpRequest request){
        if (LbRule.IPHash.name().equalsIgnoreCase(rule)){
            return request.getRemoteAddress().getAddress().getHostAddress();
        }else if(LbRule.Zone.name().equalsIgnoreCase(rule)){
            return request.getHeaders().getFirst(GatewayLoadBalanceClient.DEFAULT_ZONE_HEADER);
        }else{
            return null;
        }
    }

    private Server createServer(String info,String weight) {
        String[] infos = info.split(":");
        Assert.isTrue(infos.length==2,"server host config error,like this 127.0.0.1:80");
        WeightServer server = new WeightServer(infos[0], Integer.parseInt(infos[1]));
        if (weight!=null) {
            server.setWeight(Integer.parseInt(weight.toString()));
        }
        server.setAlive(true);
        server.setZone(GatewayLoadBalanceClient.DEFAULT_ZONE);
        return server;

    }

    private Server createServer(String info) {
        String[] infos = info.split(":");
        Assert.isTrue(infos.length==2,"server host config error,like this 127.0.0.1:80");
        Server server = null;
        try {
            server = new Server(infos[0], Integer.parseInt(infos[1]));
            server.setAlive(true);
            server.setZone(GatewayLoadBalanceClient.DEFAULT_ZONE);
        } catch (NumberFormatException e) {
            throw new IllegalArgumentException("parameter error");
        }
        return server;
    }


    @Override
    public int getOrder() {
        return 10050;
    }
}
